import os
import json
import numpy as np
import torch
from sklearn.metrics import confusion_matrix, classification_report
from pathlib import Path
import sys

# 添加项目根目录到Python路径
sys.path.append(str(Path(__file__).parent.parent.parent))

from models.factory import create_model
from data.mnist_dataset import MNISTDataModule
from core.environment import Environment
from utils.config_loader import load_config

# 所有实验目录
EXPERIMENTS = {
    "LeNet-5": os.path.join(str(Path(__file__).parent.parent.parent), "mnist_experiments_all/results_lenet5"),
    "SimpleCNN": os.path.join(str(Path(__file__).parent.parent.parent), "mnist_experiments_all/results_simple_cnn"),
    "MultiScaleCNN-Full": os.path.join(str(Path(__file__).parent.parent.parent),
                                       "mnist_experiments_all/results_multi_scale_cnn"),
    "MultiScaleCNN-NoBranch3": os.path.join(str(Path(__file__).parent.parent.parent),
                                            "mnist_experiments_all/results_multi_scale_cnn_b1-True_b2-True_b3-False")
}

# 模型名称映射（显示名 -> factory注册名）
MODEL_NAME_MAPPING = {
    "LeNet-5": "lenet5",
    "SimpleCNN": "simple_cnn",
    "MultiScaleCNN-Full": "multi_scale_cnn",
    "MultiScaleCNN-NoBranch3": "multi_scale_cnn"
}

# 输出目录
OUTPUT_DIR = Path(__file__).parent.parent.parent / "mnist_analysis"
OUTPUT_DIR.mkdir(exist_ok=True)


def load_best_model(model_name, exp_path):
    """加载最佳模型权重"""
    env_type = "local"
    config = load_config(env_type)
    env = Environment(config)

    # 关键：用映射表获取正确的工厂模型名
    factory_model_name = MODEL_NAME_MAPPING[model_name]

    # 根据模型类型创建
    if "MultiScaleCNN" in model_name:
        if "NoBranch3" in model_name:
            model = create_model(
                factory_model_name,
                num_classes=10,
                use_branch1=True,
                use_branch2=True,
                use_branch3=False
            ).to(env.device)
        else:
            model = create_model(
                factory_model_name,
                num_classes=10,
                use_branch1=True,
                use_branch2=True,
                use_branch3=True
            ).to(env.device)
    else:
        model = create_model(factory_model_name, num_classes=10).to(env.device)

    # 查找最佳模型（取第0轮）
    ckpt_dir = Path(exp_path) / "run_0" / "模型权重"
    ckpt_path = list(ckpt_dir.glob("最佳模型_epoch*.pth"))[0]
    model.load_state_dict(torch.load(ckpt_path, map_location=env.device, weights_only=True))
    model.eval()

    return model, env.device


def get_test_predictions(model, device):
    """获取测试集预测结果"""
    datamodule = MNISTDataModule(batch_size=64, data_dir=Path(__file__).parent.parent.parent / "data")
    datamodule.prepare_data()
    datamodule.setup()
    test_loader = datamodule.get_dataloaders()["val"]

    all_preds = []
    all_labels = []

    with torch.no_grad():
        for data, target in test_loader:
            data = data.to(device)
            output = model(data)
            preds = output.argmax(dim=1).cpu().numpy()
            all_preds.extend(preds)
            all_labels.extend(target.numpy())

    return np.array(all_preds), np.array(all_labels)


def main():
    for model_name, exp_path in EXPERIMENTS.items():
        print(f"正在处理: {model_name}")

        model, device = load_best_model(model_name, exp_path)
        y_pred, y_true = get_test_predictions(model, device)

        # 保存真实标签和预测标签
        np.save(OUTPUT_DIR / f"{model_name.replace(' ', '_')}_y_true.npy", y_true)
        np.save(OUTPUT_DIR / f"{model_name.replace(' ', '_')}_y_pred.npy", y_pred)

        # 混淆矩阵
        cm = confusion_matrix(y_true, y_pred)
        np.save(OUTPUT_DIR / f"{model_name.replace(' ', '_')}_confusion_matrix.npy", cm)

        # 精确率/召回率/F1
        report = classification_report(
            y_true, y_pred,
            target_names=[str(i) for i in range(10)],
            output_dict=True
        )

        with open(OUTPUT_DIR / f"{model_name.replace(' ', '_')}_metrics.json", "w", encoding="utf-8") as f:
            json.dump(report, f, ensure_ascii=False, indent=4)

        print(f"{model_name} 分析完成！")


if __name__ == "__main__":
    main()